import cv2
import os
import time
import numpy as np
from datetime import datetime

# 与 docker-compose 中 egress 转推的地址保持一致
# 注意：使用主机IP而不是容器名称，因为Python脚本运行在主机上
RTMP_URL = os.getenv("RTMP_URL", "rtmp://localhost:1935/live/room1")

# 视频处理参数
SAVE_FRAMES = os.getenv("SAVE_FRAMES", "false").lower() == "true"  # 是否保存帧
SAVE_DIR = os.getenv("SAVE_DIR", "saved_frames")                   # 保存帧的目录
DISPLAY_FPS = os.getenv("DISPLAY_FPS", "true").lower() == "true"   # 是否显示FPS
APPLY_FILTERS = os.getenv("APPLY_FILTERS", "true").lower() == "true"  # 是否应用滤镜

# 创建保存目录（如果需要）
if SAVE_FRAMES and not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)

print("连接RTMP:", RTMP_URL)
print(f"保存帧: {'是' if SAVE_FRAMES else '否'}")
print(f"显示FPS: {'是' if DISPLAY_FPS else '否'}")
print(f"应用滤镜: {'是' if APPLY_FILTERS else '否'}")

max_retries = 10
retry_count = 0
frame_count = 0
start_time = time.time()
last_saved_time = time.time()
save_interval = 5  # 每5秒保存一帧

# 可用的滤镜列表
filters = ["原始", "灰度", "边缘检测", "高斯模糊", "锐化"]
current_filter = 0

def apply_filter(frame, filter_index):
    """应用不同的图像处理滤镜"""
    if filter_index == 0:  # 原始
        return frame
    elif filter_index == 1:  # 灰度
        return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    elif filter_index == 2:  # 边缘检测
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        return cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
    elif filter_index == 3:  # 高斯模糊
        return cv2.GaussianBlur(frame, (15, 15), 0)
    elif filter_index == 4:  # 锐化
        kernel = np.array([[-1, -1, -1],
                           [-1,  9, -1],
                           [-1, -1, -1]])
        return cv2.filter2D(frame, -1, kernel)
    return frame

def save_frame(frame, directory):
    """保存当前帧到指定目录"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = os.path.join(directory, f"frame_{timestamp}.jpg")
    cv2.imwrite(filename, frame)
    print(f"已保存帧到 {filename}")

while retry_count < max_retries:
    print(f"尝试连接到RTMP流: {RTMP_URL} (尝试 {retry_count+1}/{max_retries})")
    cap = cv2.VideoCapture(RTMP_URL)
    
    if not cap.isOpened():
        print(f"连接失败，尝试重连... ({retry_count+1}/{max_retries})")
        retry_count += 1
        time.sleep(2)  # 等待2秒后重试
        continue
    
    print("连接成功！开始接收视频流")
    retry_count = 0  # 重置重试计数
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("读取帧失败，尝试重新连接...")
            break

        frame_count += 1
        current_time = time.time()
        elapsed_time = current_time - start_time
        
        # 计算FPS
        fps = frame_count / elapsed_time if elapsed_time > 0 else 0
        
        # 应用选定的滤镜
        if APPLY_FILTERS:
            processed_frame = apply_filter(frame, current_filter)
        else:
            processed_frame = frame
        
        # 在帧上添加FPS和当前滤镜信息
        if DISPLAY_FPS:
            cv2.putText(processed_frame, f"FPS: {fps:.2f}", (10, 30), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            if APPLY_FILTERS:
                cv2.putText(processed_frame, f"滤镜: {filters[current_filter]}", (10, 70), 
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 保存帧（如果启用）
        if SAVE_FRAMES and (current_time - last_saved_time) >= save_interval:
            save_frame(processed_frame, SAVE_DIR)
            last_saved_time = current_time
        
        # 显示处理后的帧
        cv2.imshow("LiveKit WebRTC Stream", processed_frame)
        
        # 键盘控制
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):  # 按q退出
            cap.release()
            cv2.destroyAllWindows()
            exit(0)
        elif key == ord('s'):  # 按s保存当前帧
            save_frame(processed_frame, SAVE_DIR)
        elif key == ord('f'):  # 按f切换滤镜
            current_filter = (current_filter + 1) % len(filters)
            print(f"已切换到滤镜: {filters[current_filter]}")

    cap.release()
    time.sleep(1)  # 重连前等待1秒

print("达到最大重试次数，退出程序")
cv2.destroyAllWindows()